# coding = utf-8

'''
展示企业数据中的分割结果可视化
'''

import os
import SimpleITK as sitk
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import cv2


def show_data(case_id, index):
    raw_root = "E:\Dataset\Qiye\DongBeiDaXue2\image_venous"
    liver_root = "E:\Dataset\Qiye\DongBeiDaXue2\liver"
    tumor_root = "F:\Dataset\Liver\qiye\DongBeiDaXue2\lesion"

    raw_root = os.path.join(raw_root, sorted(os.listdir(raw_root))[case_id - 50])
    liver_root = os.path.join(liver_root, sorted(os.listdir(liver_root))[case_id - 50])
    tumor_root = os.path.join(tumor_root, sorted(os.listdir(tumor_root))[case_id - 50])

    raw_data = sitk.GetArrayFromImage(sitk.ReadImage(raw_root))
    liver_data = sitk.GetArrayFromImage(sitk.ReadImage(liver_root))
    tumor_data = sitk.GetArrayFromImage(sitk.ReadImage(tumor_root))

    raw_data[raw_data <= -250] = -250
    raw_data[raw_data >= 200] = 200

    ours_root = "F:\predict\qiye\ours"
    hdenseunet_root = "F:\predict\qiye\hdenseunet"
    munet_root = "F:\predict\qiye\munet"
    unet_root = "F:\predict\qiye\\unet"

    ours_root = os.path.join(ours_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    hdenseunet_root = os.path.join(hdenseunet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    munet_root = os.path.join(munet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    unet_root = os.path.join(unet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))

    anchor = 0
    for i in range(liver_data.shape[0]):
        if tumor_data[i].sum() > 0 and anchor == index:
            ours_image = os.path.join(ours_root, "{}.png".format(str(index).zfill(3)))
            hdenseunet_image = os.path.join(hdenseunet_root, "{}.png".format(str(index).zfill(3)))
            munet_image = os.path.join(munet_root, "{}.png".format(str(index).zfill(3)))
            unet_image = os.path.join(unet_root, "{}.png".format(str(index).zfill(3)))

            if os.path.exists(ours_image):
                ours = Image.open(ours_image).convert("L")
                ours = np.array(ours)
                ours[ours > 0] = 1
            else:
                ours = np.zeros(tumor_data[i].shape)
            dice_ours = float(2 * (ours * tumor_data[i]).sum()) / float(ours.sum() + tumor_data[i].sum())

            if os.path.exists(hdenseunet_image):
                hdenseunet = Image.open(hdenseunet_image).convert("L")
                hdenseunet = np.array(hdenseunet)
                hdenseunet[hdenseunet > 0] = 1
            else:
                hdenseunet = np.zeros(tumor_data[i].shape)
            dice_hdenseunet = float(2 * (hdenseunet * tumor_data[i]).sum()) / float(
                hdenseunet.sum() + tumor_data[i].sum())

            if os.path.exists(munet_image):
                munet = Image.open(munet_image).convert("L")
                munet = np.array(munet)
                munet[munet > 0] = 1
            else:
                munet = np.zeros(tumor_data[i].shape)
            dice_munet = float(2 * (munet * tumor_data[i]).sum()) / float(munet.sum() + tumor_data[i].sum())

            if os.path.exists(unet_image):
                unet = Image.open(unet_image).convert("L")
                unet = np.array(unet)
                unet[unet > 0] = 1
            else:
                unet = np.zeros(tumor_data[i].shape)
            dice_unet = float(2 * (unet * tumor_data[i]).sum()) / float(unet.sum() + tumor_data[i].sum())

            print(index, round(dice_ours, 2), round(dice_hdenseunet, 2), round(dice_munet, 2), round(dice_unet, 2))

            label = np.zeros(tumor_data[i].shape)
            label[liver_data[i] == 1] = 1
            label[tumor_data[i] == 1] = 2
            label = (label/2) * 255
            label = label.astype(np.uint8)
            label = cv2.cvtColor(label, cv2.COLOR_GRAY2BGR)

            ours_lable = np.zeros(ours.shape)
            ours_lable = ours_lable.astype(np.uint8)
            ours_lable = cv2.cvtColor(ours_lable, cv2.COLOR_GRAY2BGR)
            ours_lable[ours > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if ours[x,y] > 0:
                        ours_lable[x, y] = [int(ours_lable[x,y,0]*0.5 + label[x,y,0]*0.5),
                                            int(ours_lable[x,y,1]*0.5 + label[x,y,1]*0.5),
                                            int(ours_lable[x,y,2]*0.5 + label[x,y,2]*0.5)]
                    else:
                        ours_lable[x,y] = label[x,y]

            hdenseunet_label = np.zeros(ours.shape)
            hdenseunet_label = hdenseunet_label.astype(np.uint8)
            hdenseunet_label = cv2.cvtColor(hdenseunet_label, cv2.COLOR_GRAY2BGR)
            hdenseunet_label[hdenseunet > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if hdenseunet[x, y] > 0:
                        hdenseunet_label[x, y] = [int(hdenseunet_label[x, y, 0] * 0.5 + label[x, y, 0] * 0.5),
                                            int(hdenseunet_label[x, y, 1] * 0.5 + label[x, y, 1] * 0.5),
                                            int(hdenseunet_label[x, y, 2] * 0.5 + label[x, y, 2] * 0.5)]
                    else:
                        hdenseunet_label[x, y] = label[x, y]

            munet_label = np.zeros(ours.shape)
            munet_label = munet_label.astype(np.uint8)
            munet_label = cv2.cvtColor(munet_label, cv2.COLOR_GRAY2BGR)
            munet_label[munet > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if munet[x, y] > 0:
                        munet_label[x, y] = [int(munet_label[x, y, 0] * 0.5 + label[x, y, 0] * 0.5),
                                                  int(munet_label[x, y, 1] * 0.5 + label[x, y, 1] * 0.5),
                                                  int(munet_label[x, y, 2] * 0.5 + label[x, y, 2] * 0.5)]
                    else:
                        munet_label[x, y] = label[x, y]

            unet_label = np.zeros(ours.shape)
            unet_label = unet_label.astype(np.uint8)
            unet_label = cv2.cvtColor(unet_label, cv2.COLOR_GRAY2BGR)
            unet_label[unet > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if unet[x, y] > 0:
                        unet_label[x, y] = [int(unet_label[x, y, 0] * 0.5 + label[x, y, 0] * 0.5),
                                             int(unet_label[x, y, 1] * 0.5 + label[x, y, 1] * 0.5),
                                             int(unet_label[x, y, 2] * 0.5 + label[x, y, 2] * 0.5)]
                    else:
                        unet_label[x, y] = label[x, y]

            plt.subplot(2, 3, 1)
            plt.imshow(raw_data[i], cmap="gray")
            plt.subplot(2, 3, 2)
            plt.imshow(label)
            plt.subplot(2, 3, 3)
            plt.imshow(ours_lable)
            plt.subplot(2, 3, 4)
            plt.imshow(hdenseunet_label)
            plt.subplot(2, 3, 5)
            plt.imshow(munet_label)
            plt.subplot(2, 3, 6)
            plt.imshow(unet_label)
            plt.show()

            #save
            '''
            plt.imsave("1_raw.png", raw_data[i], cmap="gray")
            plt.imsave("1_label.png", label)
            plt.imsave("1_ours.png", ours_lable)
            plt.imsave("1_hdenseunet.png", hdenseunet_label)
            plt.imsave("1_munet.png", munet_label)
            plt.imsave("1_unet.png", unet_label)
            '''


        if liver_data[i].sum() > 0:
            anchor += 1

if __name__ == '__main__':
    show_data(case_id=62, index=69)